from torch.utils.data import Dataset


class TextMatchDataset(Dataset):
    def __init__(self, x1, x2, label, num_steps,test=False):
        self.x1 = x1
        self.x2 = x2
        self.label = label
        self.num_steps = num_steps
        self.is_test = test

    def __len__(self):
        return len(self.x1)

    def __getitem__(self, idx):
        x1 = list(map(int, self.x1[idx].split()))
        x2 = list(map(int, self.x2[idx].split()))

        if len(x1) > self.num_steps:
            x1 = x1[:self.num_steps]
        if len(x1) < self.num_steps:
            x1 = x1 + [0] * (self.num_steps - len(x1))
        if len(x2) > self.num_steps:
            x2 = x2[:self.num_steps]
        if len(x2) < self.num_steps:
            x2 = x2 + [0] * (self.num_steps - len(x2))

        if self.is_test:
            return x1,x2
        else:
            return x1, x2, int(self.label[idx])
